/**
 * 
 */
package edu.ubc.javascript;

import com.google.javascript.jscomp.mozilla.rhino.ast.AstNode;
import com.google.javascript.jscomp.mozilla.rhino.ast.Block;
import com.google.javascript.jscomp.mozilla.rhino.ast.FunctionNode;
import com.google.javascript.jscomp.mozilla.rhino.ast.IfStatement;
import com.google.javascript.jscomp.mozilla.rhino.ast.Loop;
import com.google.javascript.jscomp.mozilla.rhino.ast.NodeVisitor;
import com.google.javascript.jscomp.mozilla.rhino.ast.ReturnStatement;
import com.google.javascript.jscomp.mozilla.rhino.ast.SwitchCase;

public class BlockExpansionVisitor implements NodeVisitor {

	private ReflectiveAstTransformer tx;

	public BlockExpansionVisitor(AstNode node, ReflectiveAstTransformer tx) {
		this.tx = tx;
	}
	
	@Override
	public boolean visit(AstNode node) {
		if(node instanceof IfStatement) {
			IfStatement ifStatement = (IfStatement) node;
			ifStatement.setThenPart(makeBody(ifStatement.getThenPart(), false));
			if(ifStatement.getElsePart() != null) {
				ifStatement.setElsePart(makeBody(ifStatement.getElsePart(), false));
			}
		}
		else if(node instanceof Loop) {
			Loop loop = (Loop) node;
			loop.setBody(makeBody(loop.getBody(), false));
		}
		else if(node instanceof FunctionNode) {
			FunctionNode functionNode = (FunctionNode) node;
			functionNode.setBody(makeBody(functionNode.getBody(), true));
		} else if(node instanceof SwitchCase) {
			SwitchCase switchCase = (SwitchCase) node;
			if(switchCase.getStatements() == null) {
				return true;
			}
			if(!((switchCase.getStatements().size() == 1) && (switchCase.getStatements().get(0) instanceof Block)) ) {
				Block block = new Block();
				for(AstNode caseNode : switchCase.getStatements()) {
					block.addChild(caseNode);
				}
				switchCase.getStatements().clear();
				switchCase.getStatements().add(block);
			}
		}
		return true;
	}
	
	private Block makeBody(AstNode node, boolean function) {
		if(node instanceof Block) {
			Block b = (Block) node;
			if(function) {
				b.addStatement(new ReturnStatement());
			}
			return b;
		} else {
			Block block = new Block();
			if(node != null) {
				block.addStatement(node);
				if(function) {
					block.addStatement(new ReturnStatement());
				}
			}
			return block;
		}
	}
	
}